#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
output.py - 统一支持 vibration / ae / radar
"""

import numpy as np
import requests
from datetime import datetime
import torch
import os
import zmq
import pickle
import threading
from queue import Queue
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
log = logging.getLogger(__name__)

data_queue = Queue()

# ------------------- 模型 -------------------
from vibrate.vibrate_process import process_acceleration_data
from ae.ae_process import process_AE_data
from vibrate.SAFDNN.model import SAFDNN
from ae.SAFDNN_AE.model import SAFDNN_AE

vibration_model = SAFDNN(1, 4)
if os.path.exists("./models/best.pth"):
    vibration_model.load_state_dict(torch.load("./models/best.pth", map_location='cpu'))
    vibration_model.eval()

ae_model = SAFDNN_AE(1, 2)
if os.path.exists("./models/ae_best.pth"):
    ae_model.load_state_dict(torch.load("./models/ae_best.pth", map_location='cpu'))
    ae_model.eval()


# ------------------- 雷达解析 -------------------
def parse_radar_frame(frame):
    """
    解析 0x701 帧：07 5b 82 a0 80 20 01 a0
    根据常见雷达协议（距离、角度、速度、RCS）
    """
    data = frame.get("Data", [])
    if frame.get("ID") != 0x701 or len(data) != 8:
        return None

    try:
        # Byte 0-1: 距离 (mm)
        distance = (data[0] << 8) | data[1]

        # Byte 2: 角度 (有符号，-180~180°)
        angle = data[2]
        if angle >= 128:
            angle -= 256

        # Byte 3: 相对速度 (km/h 或 m/s，需确认)
        speed = data[3]
        if speed >= 128:
            speed -= 256

        # Byte 4: RCS (dBsm)
        rcs = data[4]
        if rcs >= 128:
            rcs -= 256

        # Byte 5-7: 保留或扩展字段
        return {
            "height": distance/100,
            "angle": angle,
            "angular_velocity": speed,
            "data": data
        }
    except Exception as e:
        log.error("parse_radar_frame error: %s", e)
        return None


# ------------------- 卡信息 -------------------
def get_card_info():
    try:
        r = requests.get("http://192.168.3.88:8099/monitor/sensors/", timeout=5)
        return r.json() if r.status_code == 200 else []
    except Exception as e:
        log.error("get_card_info: %s", e)
        return []


def build_ip_dict(cards):
    d = {}
    for c in cards:
        ip = c['ip']
        if ip not in d:
            d[ip] = {'channel': [], 'type': [], 'number': []}
        d[ip]['channel'].append(c['channel'])
        d[ip]['type'].append(c['type'])
        d[ip]['number'].append(c['number'])
    return d


# ------------------- 数据处理 -------------------
def process_data():
    while True:
        client_id, msg_data = data_queue.get()
        client_ip = str(client_id, 'utf-8')

        log.info("Processing data from %s", client_ip)

        try:
            msg = pickle.loads(msg_data)
            # print(f"fmsg: {msg}")
            collect_time = msg.get('time', [])
            raw_data = msg.get('data', [])

            # 获取所有通道配置（含 radar）
            ip_dict = build_ip_dict(get_card_info())
            # print(ip_dict)
            if client_ip not in ip_dict:
                log.warning("IP %s not in config", client_ip)
                continue

            info = ip_dict[client_ip]
            if info["type"][0] == "radar":
                print(raw_data)
            buffer = np.array(raw_data, dtype=np.float64)
            idx_map = [c - min(info['channel']) for c in info['channel']]
            print(f"idx_map:{idx_map}")
            # 遍历所有通道
            for i in range(len(idx_map)):
                ch_data = buffer[idx_map[i]]
                ch_type = info['type'][i]
                ch_num = info['number'][i]
                ch_ch = info['channel'][i]
                ts = collect_time[i] if i < len(collect_time) else datetime.now().isoformat()
                print(f"ch_type:{ch_type}, ch_ch:{ch_ch}")
                if ch_type == 'vibration':
                    feats = process_acceleration_data(ch_data, 60000)
                    out = vibration_model(torch.tensor(ch_data.astype(np.float32)).unsqueeze(0).unsqueeze(0))
                    fault = torch.argmax(out, dim=1).item()
                    payload = {
                        "timeStamp": datetime.now().isoformat(),
                        "channel": ch_ch, "id": ch_num, "data": ch_data.tolist(),
                        "peak": feats[0], "peak_peak": feats[1], "mean": feats[2], "rms": feats[3],
                        "std": feats[4], "var": feats[5], "skewness": feats[6], "kurtosis": feats[7],
                        "waveform_factor": feats[8], "crest_indicator": feats[9],
                        "impulse_indicator": feats[10], "clearance_indicator": feats[11],
                        "fc": feats[12], "msf": feats[13], "vf": feats[14], "speed_rms": feats[15],
                        "GB_a_threshold": 2, "GB_v_threshold": 1.8, "fault_category": fault
                    }
                    requests.post("http://192.168.3.88:8099/monitor/vibration/sync/", json=payload, timeout=5)
                    print(f"vb done.")

                elif ch_type == 'acoustic_emission':
                    ch_data = ch_data * 20 / 8192 - 10
                    feats = process_AE_data(ch_data,1, 1000000)
                    out = ae_model(torch.tensor(ch_data.astype(np.float32)).unsqueeze(0).unsqueeze(0))
                    fault = torch.argmax(out, dim=1).item()
                    payload = {
                        "timeStamp": datetime.now().isoformat(),
                        "channel": ch_ch, "id": ch_num, "data": ch_data.tolist(),
                        "arrive_time": feats[0], "asl": feats[1], "count": feats[2],
                        "duration": feats[3], "max_amp": feats[4], "rise_time": feats[5],
                        "rms": feats[6], "fault_category": fault
                    }
                    requests.post("http://192.168.3.88:8099/monitor/acoustic/sync/", json=payload, timeout=5)
                    print(f"ae done")

                elif ch_type == 'radar':
                    print(f"msg:{msg}")
                    # 新协议：{'time': str, 'data': [[7,91,130,...]]}
                    # 旧协议：{'time': str, 'data': [...]}
                    raw_data = None

                    try:
                        data_field = msg.get('data', [])
                        print(f"data: {data_field}")
                        if not data_field:
                            log.debug("radar: data field empty")
                            continue

                        # 情况1：新协议 [ [7,91,...] ]
                        if isinstance(data_field, (list, tuple)) and len(data_field) > 0:
                            first_item = data_field[0]
                            if isinstance(first_item, (list, tuple)) and len(first_item) >= 8:
                                raw_data = list(first_item)  # 取最新一帧

                        # 情况2：旧协议 [7,91,130,...]
                        if raw_data is None and isinstance(data_field, (list, tuple)) and len(data_field) >= 8:
                            raw_data = list(data_field)

                        if raw_data is None:
                            log.debug("radar: invalid data format: %s", data_field)
                            continue

                        # 构造 frame 结构
                        frame = {
                            "ID": 0x701,
                            "Data": raw_data,
                            "DataLen": len(raw_data)
                        }

                        parsed = parse_radar_frame(frame)
                        if not parsed:
                            log.debug("radar: parse failed: %s", raw_data)
                            continue

                        payload = {
                            "timeStamp": datetime.now().isoformat(),
                            "channel": ch_ch,
                            "id": ch_num,
                            **parsed
                        }

                        try:
                            r = requests.post(
                                "http://192.168.3.88:8099/monitor/radar/sync/",
                                json=payload,
                                timeout=5
                            )
                            log.info("RADAR uploaded: %s -> %d", parsed, r.status_code)
                        except Exception as e:
                            log.error("RADAR upload failed: %s", e)

                    except Exception as e:
                        log.error("radar process error: %s | msg: %s", e, msg)

        except Exception as e:
            log.error("Process error: %s", e)
        finally:
            data_queue.task_done()


threading.Thread(target=process_data, daemon=True).start()

# ------------------- ZMQ ROUTER -------------------
context = zmq.Context()
sock = context.socket(zmq.ROUTER)
sock.bind("tcp://*:5555")
log.info("ZMQ Server started on tcp://*:5555")

while True:
    try:
        client_id, _, msg = sock.recv_multipart()
        if msg == b"HEARTBEAT":
            # print(f"receive hearbeat: {client_id}")
            sock.send_multipart([client_id, b"", b"HEARTBEAT_OK"])
            continue
        sock.send_multipart([client_id, b"", b"Array Received OK"])
        data_queue.put((client_id, msg))
    except (zmq.ZMQError, KeyboardInterrupt):
        break

context.term()